#include "intel_itt_wrapper.hpp"
#include "matrix_util.hpp"
#include "mem_usage.hpp"
#include "mpi/cipher_3d.hpp"
#include <NTL/BasicThreadPool.h>
#include <gtest/gtest.h>
#include <iostream>

#define COMM_WORLD_ROOT 0

INTEL_ITT_DOMAIN_CREATE(big_hhmm3d_domain, "test.mpi.big_hhmm3d");
INTEL_ITT_STRING_HANDLE_CREATE(hhmm3d_handle, "HMM3D");

const double MB = static_cast<double>(1l << 20l);

params<helib::BGV> my_params;
bool no_encoding_reuse = false;
bool local_check = false;
bool no_nested = false;
bool check = false;
bool shared_4d_lagre = false;
bool shared_4d = false;
bool fhe4d = false;
bool distributed_4d = false;
bool warming = false;
unsigned long input_bit_len;
unsigned long hmm_threads = 0;
unsigned long mp = 0;
unsigned long np = 0;
unsigned long kp = 0;

#define PCGEMM_3D_TEST(DIM_M, DIM_N, DIM_K, OPT)                               \
  TEST(TEST_PCGEMM_3D_OPT##OPT, M##DIM_M##_N##DIM_N##_K##DIM_K) {              \
    int rank;                                                                  \
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);                                      \
                                                                               \
    if (local_check) {                                                         \
      check = local_check;                                                     \
    }                                                                          \
    std::size_t mem_begin_ = 0.0;                                              \
    std::size_t mem_end_ = 0.0;                                                \
                                                                               \
    unsigned long m = DIM_M;                                                   \
    unsigned long n = DIM_N;                                                   \
    unsigned long k = DIM_K;                                                   \
    hmm_status<shmm_engine> status(METHOD::BSGS);                              \
    if (shared_4d) {                                                           \
      distributed_4d = false;                                                  \
    }                                                                          \
    if (distributed_4d) {                                                      \
      fhe4d = true;                                                            \
      shared_4d = false;                                                       \
    }                                                                          \
    mhhmm_3d ce(my_params, status, COMM_WORLD_ROOT, MPI_COMM_WORLD, m, n, k,   \
                OPT, fhe4d, mp, np, kp);                                       \
                                                                               \
    int *A = nullptr;                                                          \
    int *B = nullptr;                                                          \
    int *C = nullptr;                                                          \
    if (rank == COMM_WORLD_ROOT) {                                             \
      const std::size_t size_A = m * k;                                        \
      const std::size_t size_B = k * n;                                        \
      FIX_RAND_SEED(true);                                                     \
      ALLOC_INT_ARRAY_WITHOUT_INIT(A, size_A);                                 \
      ALLOC_INT_ARRAY_WITHOUT_INIT(B, size_B);                                 \
      RAND_ARRAY(A, size_A, input_bit_len);                                    \
      RAND_ARRAY(B, size_B, input_bit_len);                                    \
      /* PRINTF_ARRAY_AS_MATRIX(std::cout, A, m, k);                           \
      PRINTF_ARRAY_AS_MATRIX(std::cout, B, k, n);*/                            \
      if (check) {                                                             \
        const std::size_t size_C = m * n;                                      \
        ALLOC_INT_ARRAY_WITHOUT_INIT(C, size_C);                               \
        cblas_igemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0,   \
                    A, k, B, n, 0.0, C, n);                                    \
      }                                                                        \
    }                                                                          \
                                                                               \
    /* Distribute Enc(input) */                                                \
    if (rank == COMM_WORLD_ROOT) {                                             \
      std::cout << "Init data - Start: " << MPI_Wtime() << " s\n";             \
      std::flush(std::cout);                                                   \
    }                                                                          \
    ce.init_data(A, B);                                                        \
    if (rank == COMM_WORLD_ROOT) {                                             \
      std::cout << "Init data - Stop: " << MPI_Wtime() << " s\n";              \
      std::flush(std::cout);                                                   \
    }                                                                          \
                                                                               \
    /* Execute Cannon's algorithm & Get partial Enc(result) */                 \
    ce.hmme_status()->threads(hmm_threads);                                    \
    ce.set_nest(!no_nested);                                                   \
    if (no_encoding_reuse) {                                                   \
      ce.hmme_status()->reset(false);                                          \
      ce.hmme_status()->set_encoding_reuse(false);                             \
    }                                                                          \
    if (rank == COMM_WORLD_ROOT) {                                             \
      std::stringstream ss;                                                    \
      ss << "* OMP Nested            : " << ce.nest() << std::endl;            \
      ss << "* Encoding Reuse        : " << ce.hmme_status()->encoding_reuse() \
         << std::endl;                                                         \
      ss << "* CGEMM Method          : "                                       \
         << shmm_engine::METHOD_STR[ce.hmme_status()->method()] << std::endl;  \
      auto threads = ce.hmme_status()->threads();                              \
      ss << "* CGEMM # of threads    : "                                       \
         << (threads == 0 ? "omptask" : std::to_string(threads)) << std::endl; \
      std::cout << ss.str();                                                   \
      std::flush(std::cout);                                                   \
    }                                                                          \
    std::vector<std::string> partial_C;                                        \
    if (warming) {                                                             \
      ce.exec(partial_C);                                                      \
      ce.clear();                                                              \
      partial_C.clear();                                                       \
    }                                                                          \
    if (rank == COMM_WORLD_ROOT) {                                             \
      std::cout << "Execution - Start: " << MPI_Wtime() << " s\n";             \
      std::flush(std::cout);                                                   \
    }                                                                          \
    INTEL_ITT_RESUME;                                                          \
    INTEL_ITT_TASK_BEGIN(big_hhmm3d_domain, hhmm3d_handle);                    \
    mem_begin_ = getCurrentRSS();                                              \
    if (shared_4d) {                                                           \
      ce.exec_4d_shared(partial_C, shared_4d_lagre);                           \
    } else {                                                                   \
      ce.exec(partial_C, distributed_4d);                                      \
    }                                                                          \
    mem_end_ = getPeakRSS();                                                   \
    auto mem_algor = static_cast<double>(mem_end_ - mem_begin_) / MB;          \
    INTEL_ITT_TASK_END(big_hhmm3d_domain);                                     \
    INTEL_ITT_DETACH;                                                          \
    if (rank == COMM_WORLD_ROOT) {                                             \
      std::cout << "Execution - Stop: " << MPI_Wtime() << " s\n";              \
      std::cout << "HMM Mem Usage   : " << mem_algor                           \
                << " MB = " << mem_algor / 1024.0 << " GB\n";                  \
      std::flush(std::cout);                                                   \
    }                                                                          \
    if (rank == COMM_WORLD_ROOT) {                                             \
      std::stringstream ss;                                                    \
      ce.report(ss);                                                           \
      std::cout << ss.str();                                                   \
      std::flush(std::cout);                                                   \
    }                                                                          \
                                                                               \
    std::vector<NTL::mat_ZZ> res;                                              \
    if (check) {                                                               \
      /* Only the root will get a vector of Dec(Enc(result)) */                \
      if (rank == COMM_WORLD_ROOT) {                                           \
        std::cout << "Collect Data - Start : " << MPI_Wtime() << " s\n";       \
      }                                                                        \
      ce.collect_data(partial_C, res, local_check);                            \
      if (rank == COMM_WORLD_ROOT) {                                           \
        std::cout << "Collect Data - Stop  : " << MPI_Wtime() << " s\n";       \
        std::flush(std::cout);                                                 \
      }                                                                        \
    }                                                                          \
                                                                               \
    if (rank != COMM_WORLD_ROOT) {                                             \
      return;                                                                  \
    }                                                                          \
    if (!check) {                                                              \
      delete[] A;                                                              \
      delete[] B;                                                              \
      GTEST_SKIP() << "Skip checking answers";                                 \
    }                                                                          \
    /* Check answers */                                                        \
    std::cout << "Check answers - Start: " << MPI_Wtime() << " s\n";           \
    const auto mb = ce.block_m();                                              \
    const auto nb = ce.block_n();                                              \
    const auto mp = ce.mp();                                                   \
    const auto np = ce.np();                                                   \
    long hits = 0;                                                             \
    for (std::size_t i = 0; i < mp; i++) {                                     \
      for (std::size_t j = 0; j < np; j++) {                                   \
        if (local_check && (i > 0 || j > 0)) {                                 \
          break;                                                               \
        }                                                                      \
        for (std::size_t ii = 0; ii < mb; ii++) {                              \
          for (std::size_t jj = 0; jj < nb; jj++) {                            \
            auto expect = C[(i * mb + ii) * n + j * nb + jj];                  \
            auto actual = NTL::conv<int>(res[i * np + j][ii][jj]);             \
            ASSERT_EQ(expect, actual)                                          \
                << "diff at (" << ii << ", " << jj << ") in BLK (" << i        \
                << ", " << j << ")";                                           \
            hits++;                                                            \
          }                                                                    \
        }                                                                      \
      }                                                                        \
    }                                                                          \
    std::cout << "Check answers - Stop : " << MPI_Wtime() << " s\n";           \
    std::cout << "Check answers        : " << hits << " passed\n";             \
    std::flush(std::cout);                                                     \
    delete[] A;                                                                \
    delete[] B;                                                                \
    delete[] C;                                                                \
  }

/* Local Check */
// PCGEMM_3D_TEST(16, 16, 16, 2)
PCGEMM_3D_TEST(32, 32, 4, 4)
PCGEMM_3D_TEST(32, 32, 32, 4)

/* Medium Check */
PCGEMM_3D_TEST(512, 512, 16, 128)
PCGEMM_3D_TEST(4096, 64, 32, 128)
PCGEMM_3D_TEST(512, 512, 512, 128)
PCGEMM_3D_TEST(4096, 16, 4096, 128)
PCGEMM_3D_TEST(16, 4096, 4096, 128)
PCGEMM_3D_TEST(4096, 4096, 16, 128)
PCGEMM_3D_TEST(1024, 1024, 1024, 128)
PCGEMM_3D_TEST(32768, 64, 16, 128)
PCGEMM_3D_TEST(2048, 2048, 256, 128)
/* TDSC2022 */
PCGEMM_3D_TEST(2048, 2048, 8, 128)
PCGEMM_3D_TEST(1024, 1024, 16, 128)
PCGEMM_3D_TEST(1024, 16, 1024, 128)

/* Large-scale Tests */
PCGEMM_3D_TEST(2048, 2048, 2048, 128)
PCGEMM_3D_TEST(16384, 16, 16384, 128)
PCGEMM_3D_TEST(1048576, 64, 16, 128)
PCGEMM_3D_TEST(1048576, 16, 64, 128)
PCGEMM_3D_TEST(16384, 16384, 16, 128)

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  int ret = MPI_Init(&argc, &argv);
  EXPECT_EQ(ret, 0);

  INTEL_ITT_PAUSE;
  // Get MPI metadata
  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  // Only one printer
  ::testing::TestEventListeners &listeners =
      ::testing::UnitTest::GetInstance()->listeners();
  if (rank != COMM_WORLD_ROOT) {
    //* Only the root has a secret key
    delete listeners.Release(listeners.default_result_printer());
  }

  unsigned long ntl_threads;
  auto max_omp_threads = omp_get_max_threads();

  // Parse command line arguments
  command_line_parser parser(my_params);
  parser.add_option("-nn", "no_omp_nested", &no_nested);
  parser.add_option("-chk", "enable_check_answer", &check);
  parser.add_option("-lchk", "enable_local_check", &local_check);
  parser.add_option("-ner", "no_encoding_reuse", &no_encoding_reuse);
  parser.add_option("-4dp", "enable_fhe4d", &fhe4d);
  parser.add_option("-lm", "enable_shared_4d_large", &shared_4d_lagre);
  parser.add_option("-lo", "enable_shared_4d", &shared_4d);
  parser.add_option("-4d", "enable_distributed_4d", &distributed_4d);
  parser.add_option("-wu", "enable_warming_up", &warming);
  parser.add_option("-ib", "7", "input_bit_length", &input_bit_len);
  parser.add_option("-ht", std::to_string(max_omp_threads), "hmm_threads",
                    &hmm_threads);
  parser.add_option("-nt", "1", "ntl_threads", &ntl_threads);
  parser.add_option("-mp", "0", "mp", &mp);
  parser.add_option("-np", "0", "np", &np);
  parser.add_option("-kp", "0", "kp", &kp);
  ret = parser.parse(argc, argv, rank == COMM_WORLD_ROOT);
  EXPECT_EQ(ret, 0);

  NTL::SetNumThreads(
      std::min(ntl_threads, static_cast<ulong>(max_omp_threads)));

  // Run tests
  ret = RUN_ALL_TESTS();
  EXPECT_EQ(ret, 0);

  MPI_Finalize();
  return 0;
}
